package com.ruoyi.ai.bootstrap;

import org.springframework.ai.reader.tika.TikaDocumentReader;
import com.ruoyi.ai.service.VectorService;
import com.ruoyi.ai.splitter.LineTextSplitter;
import com.ruoyi.ai.splitter.MarkdownSplitter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.transformer.splitter.TextSplitter;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.boot.CommandLineRunner;
import org.springframework.core.io.ClassPathResource;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

/**
 * 向量数据库初始化加载器
 * 该类负责在应用启动时初始化示例商品数据到向量数据库中
 * 实现CommandLineRunner接口以确保在Spring Boot应用启动时自动执行初始化逻辑
 *
 * @author tanxiaolong
 */
@Component
public class LoadVectorStore implements CommandLineRunner {
    /**
     * 日志记录器
     */
    private static final Logger logger = LoggerFactory.getLogger(LoadVectorStore.class);

    /**
     * 向量服务接口，用于商品数据的存储和检索
     */
    @Autowired
    private  VectorService vectorService;

    @Autowired
    @Qualifier("milvusVectorStore")
    private  VectorStore vectorStore;


    /**
     * Spring Boot应用启动时自动执行的初始化方法
     * 检查向量数据库中是否存在商品数据，如果为空则加载预设的示例商品数据
     *
     * @param args 启动参数
     * @throws Exception 初始化过程中可能发生的异常
     */
    @Override
    public void run(String... args) throws Exception {
        logger.debug("开始初始化问题数据");
        if (args.length > 0) {
            logger.info("接收到参数: {}", args[0]);
        }

        if (Objects.requireNonNull(vectorService.searchQuestion("什么是重载，重载的理解，重载的优点，重载有什么用，重载的体系介绍")).isEmpty()) {
            logger.debug("问题数据为空，重新加载");
            List<Resource> list = List.of(
                    new ClassPathResource("/doc/1-Java基础重要知识扫盲.md")
            );
            //读取文件
            list.forEach(document -> {
                TikaDocumentReader tikaDocumentReader = new TikaDocumentReader(document);
                List<Document> documents = tikaDocumentReader.get();
                TextSplitter textSplitter = null;
                //判断Resource文档类型
                if (Objects.requireNonNull(document.getFilename()).endsWith(".md")) {
                    //使用MarkdownSplitter切分
                    textSplitter = new MarkdownSplitter();
                } else if (document.getFilename().endsWith(".txt")) {
                    textSplitter = new LineTextSplitter();
                } else {
                    //使用TokenTextSplitter切分
                    textSplitter = new TokenTextSplitter();
                }
                List<Document> splitDocuments = textSplitter.split(documents.get(0));
                // 分批处理文档，每批最多10个
                List<List<Document>> batches = new ArrayList<>();
                for (int i = 0; i < splitDocuments.size(); i += 10) {
                    int end = Math.min(i + 10, splitDocuments.size());
                    batches.add(splitDocuments.subList(i, end));
                }

                // 逐批添加到向量存储
                for (List<Document> batch : batches) {
                    vectorStore.add(batch);
                }
                logger.info("初始化了{}个示例问题数据", splitDocuments.size());
            });

        }
        logger.debug("初始化数据完毕");
    }
}
